import os

from manimlib import *
from collections import deque


class ClickBlock(Button):

    def __init__(self, side_length, index=0, scene=None, **kwargs):
        assert scene
        mobject = VGroup(
            Square(side_length, **kwargs),
            Integer(index)
        )
        self.index = index
        self.scene_ctx = scene
        super().__init__(mobject, self.on_click, **kwargs)
        self.data["can_move"] = ORIGIN

    def on_click(self, mob):
        if np.all(self.data["can_move"] == ORIGIN):
            return
        scene = self.scene_ctx
        scene.play(mob.animate.shift(scene.shift_distance * self.data["can_move"]),
                   run_time=0.4, rate_func=rush_from)
        scene.swap_array_item(self.index, 0)
        scene.update_can_move()


class EightPuzzleScene(Scene):
    CONFIG = {
        "square_side_length": 2,
        "square_buff": 0.2,
        "box_center": LEFT * 2.5,
        "array": np.array([1, 2, 3, 4, 5, 6, 7, 8, 0]),
        "color_list": [RED, ORANGE, GOLD, YELLOW, GREEN, TEAL, BLUE, PURPLE],
        "on_anim": False
    }

    def init_game_box(self):
        box = Square(3 * self.square_side_length + 4 * self.square_buff)
        box.shift(self.box_center)
        self.add(box)
        self.shift_distance = self.square_side_length + self.square_buff

    def init_squares(self):
        self.squares = Group()
        for i in range(len(self.array) - 1):
            self.squares.add(ClickBlock(
                self.square_side_length,
                index=self.array[i],
                scene=self,
                color=self.color_list[i])
            )
        self.squares.arrange_in_grid(3, 3, self.square_buff)
        self.squares.move_to(self.box_center)
        self.add(self.squares)
        # self.update_can_move()     # this line may cause a bug

    def init_buttons(self):
        self.info = Text("单击正方形 将它挪到空格处", font="Source Han Serif SC").add_background_rectangle(
            color=BLACK, buff=0.25)
        self.info.to_edge(DOWN)

        play_rec = Text("Play").add_background_rectangle(
            color=BLACK, buff=0.25)
        play_btn = Button(play_rec, self.start_to_play)
        info_rec = Text("Help").add_background_rectangle(
            color=BLACK, buff=0.25)
        info_btn = Button(info_rec, self.show_info)
        solve_rec = Text("Solve").add_background_rectangle(
            color=BLACK, buff=0.25)
        solve_btn = Button(solve_rec, self.solve_problem)
        play_btn.shift(UP * 1.2 + RIGHT * 5)
        info_btn.shift(DOWN * 1.2 + RIGHT * 5)
        solve_btn.shift(RIGHT * 5)
        self.add(play_btn, info_btn, solve_btn)

    def update_can_move(self):
        zero_index = np.where(self.array == 0)[0][0]
        for i in range(len(self.array)):
            if zero_index == i:
                continue
            el = self.array[i]  # the element of array that index==i
            if i == zero_index - 1:
                self.squares[el - 1].data["can_move"] = RIGHT
            elif i == zero_index + 1:
                self.squares[el - 1].data["can_move"] = LEFT
            elif i == zero_index - 3:
                self.squares[el - 1].data["can_move"] = DOWN
            elif i == zero_index + 3:
                self.squares[el - 1].data["can_move"] = UP
            else:
                self.squares[el - 1].data["can_move"] = ORIGIN

    def start_to_play(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        random.seed(time.time())
        random.shuffle(self.array)
        while self.sln.has_answer(self.array) == False:
            random.shuffle(self.array)

        pos = dict()
        relative_center = self.box_center + UL * self.shift_distance
        for i in range(len(self.array)):  # i is the index of array
            if self.array[i] == 0:
                continue
            index = self.array[i] - 1  # the index of square
            row, col = i // 3, i % 3
            pos[index] = relative_center + row * DOWN * \
                         self.shift_distance + col * RIGHT * self.shift_distance
        self.play(LaggedStart(*[
            ApplyMethod(self.squares[i].move_to, pos[i],
                        path_arc=PI * (random.random() - 0.5))
            for i in pos
        ]), run_time=1.5, lag_ratio=0.1)
        self.update_can_move()
        self.on_anim = False

    def solve_problem(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        print("Solving...")
        ima = time.time()
        steps = self.sln.get_answer2(self.array)
        print(steps)
        print("time =", time.time() - ima)
        if steps is None:
            self.on_anim = False
            return
        for i in steps:
            index = self.array[i] - 1
            self.squares[index].on_click(self.squares[index])
        print("Finish!")
        self.on_anim = False

    def show_info(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        self.play(FadeIn(self.info, UP))
        self.wait()
        self.play(FadeOut(self.info, DOWN))
        self.on_anim = False

    def swap_array_item(self, i, j):
        a = np.where(self.array == i)[0][0]
        b = np.where(self.array == j)[0][0]
        # a, b = np.where(self.array == i or self.array == j)[0]
        self.array[a], self.array[b] = self.array[b], self.array[a]

    def setup(self):
        self.sln = Solution()
        self.init_game_box()
        self.init_squares()
        self.init_buttons()


class Solution:
    """
    传入一个字符串, 返回一个操作序列, 
    这个序列是绝对坐标的点击序列
    """

    def __init__(self):
        self.answer = np.array([1, 2, 3, 4, 5, 6, 7, 8, 0])

    def get_array_that_ai_larger_than_previous_numbers(self, array: np.ndarray) -> np.ndarray:
        res = np.zeros(8)
        mask = array != 0
        sub_arr = array[mask]
        for i in range(1, len(sub_arr)):
            query = np.where(sub_arr[:i] < sub_arr[i])[0]
            res[i] = len(query)
        return res

    def has_answer(self, start: np.ndarray) -> bool:
        sum_array = sum(self.get_array_that_ai_larger_than_previous_numbers(start))
        return sum_array % 2 == 0

    def get_answer(self, start: np.ndarray) -> list:
        """
        To slow!
        """
        d = {start.tobytes(): []}
        q = deque()
        q.append(start)
        dx = [0, 0, 1, -1]
        dy = [1, -1, 0, 0]
        while q:
            t = q.popleft()
            steps = d[t.tobytes()].copy()
            if np.all(t == self.answer):
                return steps
            k = np.where(t == 0)[0][0]
            x, y = k // 3, k % 3

            for i in range(4):
                a = x + dx[i]
                b = y + dy[i]
                if 0 <= a < 3 and 0 <= b < 3:
                    nk = a * 3 + b
                    t[nk], t[k] = t[k], t[nk]
                    if t.tobytes() not in d:
                        q.append(t.copy())
                        d[t.tobytes()] = steps + [nk]
                    t[nk], t[k] = t[k], t[nk]
        return None

    def get_string_array(self, arr: np.ndarray) -> str:
        return ''.join(list(map(str, arr)))

    def get_answer2(self, arr: np.ndarray):
        start = self.get_string_array(arr)
        d = {start: []}
        q = deque()
        q.append(start)
        dx = [0, 0, 1, -1]
        dy = [1, -1, 0, 0]
        result = list()
        while q:
            t = q.popleft()
            distance = d[t].copy()
            if t == "123456780":
                return d[t]
            k = t.index('0')
            x, y = k // 3, k % 3

            tl = list(t)
            for i in range(4):
                a = x + dx[i]
                b = y + dy[i]
                if 0 <= a < 3 and 0 <= b < 3:
                    nk = a * 3 + b
                    tl[nk], tl[k] = tl[k], tl[nk]
                    t = "".join(tl)
                    if t not in d:
                        q.append(t)
                        d[t] = distance + [nk]
                    tl[nk], tl[k] = tl[k], tl[nk]
        return None
